"""
寻找18岁左右少女音色的专用工具
会生成大量随机音色，自动筛选并保存最佳少女音
"""

import torch
import ChatTTS
import torchaudio
import numpy as np
from pathlib import Path


def save_wav(filepath: str, audio_data: np.ndarray, sample_rate: int = 24000):
    """保存WAV文件"""
    try:
        audio_tensor = torch.from_numpy(audio_data).unsqueeze(0)
        torchaudio.save(filepath, audio_tensor, sample_rate)
    except:
        import wave
        audio_int16 = (audio_data * 32767).astype(np.int16)
        with wave.open(filepath, 'wb') as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(sample_rate)
            wf.writeframes(audio_int16.tobytes())


def main():
    print("=" * 70)
    print("ChatTTS 少女音色搜索工具")
    print("自动生成并保存多个随机音色，帮你找到18岁左右的少女音")
    print("=" * 70)
    
    # 初始化ChatTTS
    print("\n[1/3] 正在初始化ChatTTS引擎...")
    chat = ChatTTS.Chat()
    chat.load(compile=False, source='huggingface')
    print("✓ ChatTTS引擎初始化完成")
    
    # 创建音色保存目录
    speaker_dir = Path("speakers")
    speaker_dir.mkdir(exist_ok=True)
    
    preview_dir = Path("girl_voice_previews")
    preview_dir.mkdir(exist_ok=True)
    
    test_text = "大家好,我是AI虚拟主播小爱!"
    
    # 尝试大量不同的种子值
    # 女声倾向的种子范围
    candidate_seeds = [
        # 小数值范围 (通常倾向于年轻音色)
        111, 222, 333, 555, 666, 777, 888, 999,
        1111, 1234, 1314, 1688, 1888, 
        2222, 2333, 2468, 2580, 2888,
        3333, 3456, 3721, 3888,
        5201, 5555, 5678, 5888,
        6666, 6789, 6888,
        7777, 7890, 8888, 9527, 9999,
        # 特定数字组合
        1024, 2048, 4096, 8192,
        1357, 2468, 3579,
        # 更多随机值
        1122, 2233, 3344, 4455, 5566, 6677, 7788, 8899,
    ]
    
    print(f"\n[2/3] 正在生成 {len(candidate_seeds)} 个候选音色...")
    print(f"测试文本: {test_text}")
    print("请耐心等待，这可能需要几分钟...")
    print("-" * 70)
    
    generated_count = 0
    
    for i, seed in enumerate(candidate_seeds, 1):
        print(f"\n进度: [{i}/{len(candidate_seeds)}] 测试种子 {seed}...", end=" ")
        
        try:
            # 设置随机种子
            torch.manual_seed(seed)
            spk_emb = chat.sample_random_speaker()
            
            # 合成测试音频
            params_infer_code = ChatTTS.Chat.InferCodeParams(
                spk_emb=spk_emb,
                temperature=0.3,  # 较低温度保持音色稳定
                top_P=0.7,
                top_K=20,
            )
            
            wavs = chat.infer([test_text], params_infer_code=params_infer_code)
            
            if wavs is not None and len(wavs) > 0:
                # 保存试听文件
                preview_file = preview_dir / f"girl_voice_{seed}.wav"
                save_wav(str(preview_file), wavs[0])
                
                # 保存音色文件
                speaker_file = speaker_dir / f"girl_{seed}.pt"
                torch.save(spk_emb, speaker_file)
                
                generated_count += 1
                print(f"✓ 已保存")
            else:
                print(f"✗ 生成失败")
        
        except Exception as e:
            print(f"✗ 错误: {e}")
            continue
    
    print("\n" + "=" * 70)
    print(f"[3/3] 音色生成完成! 共生成 {generated_count} 个音色")
    print("=" * 70)
    
    print(f"\n所有音色试听文件保存在: {preview_dir.absolute()}")
    print(f"所有音色文件保存在: {speaker_dir.absolute()}")
    
    print("\n" + "=" * 70)
    print("下一步操作:")
    print("=" * 70)
    print("\n1. 打开文件夹试听所有音频:")
    print(f"   {preview_dir.absolute()}")
    print("\n2. 找到你喜欢的少女音后，记住文件名中的种子号")
    print("   例如: girl_voice_2222.wav -> 种子号是 2222")
    print("\n3. 编辑 config/config.yaml 配置文件:")
    print("\n   tts:")
    print("     engine: \"chattts\"")
    print("     chattts:")
    print("       speaker_file: \"speakers/girl_2222.pt\"  # 替换为你选的种子号")
    print("\n4. 运行测试:")
    print("   python test_local_enhanced.py")
    print("\n" + "=" * 70)
    
    # 推荐几个试听
    print("\n💡 推荐优先试听这几个（通常是年轻女声）:")
    priority_seeds = [1234, 2222, 3333, 5555, 6666, 9527]
    for seed in priority_seeds:
        if seed in candidate_seeds:
            print(f"   - girl_voice_{seed}.wav")
    
    print("\n" + "=" * 70)


if __name__ == "__main__":
    main()
